import sensor, image, time, math
import pyb
from pyb import UART
import ulab as np
import gc

# LED状态指示初始化
pyb.LED(1).on()   # 红色LED指示电源
pyb.LED(2).off()  # 绿色LED指示摄像头状态
pyb.LED(3).off()  # 蓝色LED指示通信状态

# #################### 常量定义 ####################
# 系统操作模式字典（状态机基础）
OPERATION_MODE = {
    'STANDBY': 0,                 # 待机模式
    'LINE_TRACKING': 1,           # 巡线模式
    'COLOR_DETECTION': 2,         # 物料颜色识别
    'COARSE_AREA_DETECTION': 3,   # 粗加工区定位
    'BUFFER_AREA_DETECTION': 4,   # 缓存区定位
    'QR_CODE_SCAN': 5             # 二维码扫描
}

# UART通信常量
UART_HEADER = 0xAA
UART_FOOTER = 0xBB

# 命令映射表 (主命令)
COMMAND_MAP = {
    0x01: 'LINE_TRACKING',
    0x02: 'COLOR_DETECTION',
    0x03: 'COARSE_AREA_DETECTION',
    0x04: 'BUFFER_AREA_DETECTION',
    0x05: 'QR_CODE_SCAN',
    0x00: 'STANDBY'
}

# 颜色子命令映射 (用于颜色检测模式)
COLOR_COMMAND_MAP = {
    0x01: 'red',
    0x02: 'green',
    0x03: 'blue'
}

# 优化后的颜色阈值 - LAB色彩空间
COLOR_THRESHOLDS = {
    'red': [(17, 32, 3, 37, 2, 33)],          # 优化红色阈值
    'green': [(64, 75, -63, -20, 8, 50)],       # 优化绿色阈值
    'blue': [(30, 67, -3, 10, -54, -25)],      # 统一蓝色阈值
    'line_tracking': [(0, 30, 0, 64, -128, -20)],  # 黑线阈值
    # 'coarse_red': [(17, 56, 3, 13, -3, 21)],
    # 'coarse_green': [(23, 57, -36, -19, -9, 19)],
    # 'coarse_blue': [(22, 62, -22, -11, -19, -1)]
}



# 颜色识别结果代码 (与命令代码区分开)
COLOR_RESULT_CODES = {
    'red': 0x1,
    'green': 0x2,
    'blue': 0x3
}

DETECTION_PARAMS = {
    'line_regression': {
        'x_stride': 2,
        'y_stride': 2,
        'threshold': 1000,
        'theta_margin': 25,
        'rho_margin': 25
    },
    'coarse_area_circle': {
        'r_min': 30,
        'r_max': 95,
        'samples_needed': 6,  # 增加采样点
        'area_threshold': 3000  # 降低阈值，增加识别率
    }
}

# #################### 初始化设置 ####################
# 摄像头初始化
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA)    # 320x240分辨率
sensor.skip_frames(time=2000)
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(False)
sensor.set_auto_exposure(False, exposure_us=8000)  # 适当增加曝光时间
sensor.set_contrast(1)
# sensor.set_brightness(0)             # 设置亮度
# sensor.set_saturation(1)             # 增加饱和度以便更好地区分颜色

# UART初始化
uart = UART(3, 115200)
uart.init(115200, bits=8, parity=None, stop=1)

# 全局变量
current_operation = OPERATION_MODE['STANDBY']
target_color = None  # 目标检测颜色
line_x_samples = [0] * 15
line_y_samples = [0] * 15
sample_counter = 0
# DEBUG_MODE = True  # 调试模式开关
# DEBUG_TARGET_COLOR = 'blue'


# #################### 功能函数 ####################
def send_data(*payload):
    """发送数据包到STM32（优化版）"""
    # 仅清空发送缓冲区，保留接收数据
    uart.writechar(UART_HEADER)
    for byte in payload:
        uart.writechar(byte & 0xFF)
    uart.writechar(UART_FOOTER)
    pyb.delay(2)


def parse_command():
    """解析从STM32接收的命令（增强版）"""
    if uart.any():
        raw_data = uart.read()
        data_bytes = bytearray(raw_data)

        # 调试输出原始数据
        #print("Raw RX:", [hex(b) for b in data_bytes])

        # 查找有效数据包起始位置
        start_idx = -1
        for i in range(len(data_bytes)):
            if data_bytes[i] == UART_HEADER:
                start_idx = i
                break

        # 未找到帧头则直接返回
        if start_idx == -1:
            return {'valid': False}

        # 截取有效数据段
        data_bytes = data_bytes[start_idx:]

        # 查找帧尾位置
        end_idx = -1
        for i in range(len(data_bytes)):
            if data_bytes[i] == UART_FOOTER:
                end_idx = i
                break

        # 未找到完整数据包
        if end_idx == -1 or end_idx <= start_idx:
            return {'valid': False}

        # 提取完整数据包
        packet = data_bytes[:end_idx+1]
        print("Processing packet:", [hex(b) for b in packet])

        # 校验数据包长度
        if len(packet) < 4:  # 最小有效包：AA CMD DATA BB
            return {'valid': False}

        # 解析命令
        result = {
            'valid': False,
            'command': 'UNKNOWN',
            'color': None
        }

        cmd_byte = packet[1]
        if cmd_byte in COMMAND_MAP:
            result['command'] = COMMAND_MAP[cmd_byte]
            result['valid'] = True

            # 处理颜色检测子命令
            if result['command'] == 'COLOR_DETECTION' and len(packet) >=4:
                color_cmd = packet[2]
                if color_cmd in COLOR_COMMAND_MAP:
                    result['color'] = COLOR_COMMAND_MAP[color_cmd]


            # 处理包含颜色子命令的情况（支持颜色检测和粗加工区检测）
                    if result['command'] in ['COLOR_DETECTION', 'COARSE_AREA_DETECTION'] and len(packet) >=4:
                        color_cmd = packet[2]
                        if color_cmd in COLOR_COMMAND_MAP:
                            result['color'] = COLOR_COMMAND_MAP[color_cmd]

        return result

    return {'valid': False}


def detect_specific_color(img, target_color):
    """检测指定颜色的物体，返回检测结果"""
    # 定义检测参数
    min_blob_area = 5000      # 最小面积要求
    roi = (20, 20, 120, 80)  # ROI区域

    # 在ROI区域绘制矩形框(调试用)
    img.draw_rectangle(roi, color=(255,255,255), thickness=1)

    # 获取目标颜色的阈值
    threshold = COLOR_THRESHOLDS.get(target_color)
    if not threshold:
        print(f"未知颜色: {target_color}")
        return None

    # 根据颜色选择标记颜色
    if target_color == 'red':
        mark_color = (255, 0, 0)
    elif target_color == 'green':
        mark_color = (0, 255, 0)
    else:  # 蓝色
        mark_color = (0, 0, 255)

    # 查找色块
    blobs = img.find_blobs(
        threshold,
        roi=roi,
        pixels_threshold=500,    # 像素阈值
        area_threshold=min_blob_area,
        merge=True,
        margin=10
    )

    # 找到有效的色块
    if blobs:
        # 选择最大的色块
        largest_blob = max(blobs, key=lambda b: b.area())

        # 检查形状是否合理(排除长条或奇怪形状)
        ratio = largest_blob.w() / largest_blob.h() if largest_blob.h() != 0 else 0
        if 0.5 <= ratio <= 2.0:  # 限制宽高比
            # 绘制检测结果
            img.draw_rectangle(largest_blob.rect(), color=mark_color, thickness=2)
            img.draw_cross(largest_blob.cx(), largest_blob.cy(), color=mark_color, thickness=2)
            img.draw_string(largest_blob.cx()-10, largest_blob.cy()-30, target_color.upper(),
                           color=mark_color, scale=2)

            # 计算归一化坐标(0-255范围)
            norm_x = int(largest_blob.cx() * 255 / img.width())
            norm_y = int(largest_blob.cy() * 255 / img.height())

            # 返回检测结果
            return {
                'color': target_color,
                'code': COLOR_RESULT_CODES[target_color],  # 使用结果代码，而非命令代码
                'x': norm_x,
                'y': norm_y,
                'area': largest_blob.area()
            }

    return None

# #################### 主循环 ####################
clock = time.clock()


# 主循环
while True:
    clock.tick()
    img = sensor.snapshot()

    # 清理内存
    gc.collect()


    # DEBUG_MODE = True  # 调试开关
    # if DEBUG_MODE:
    #     current_operation = OPERATION_MODE['COLOR_DETECTION']
    #     target_color = 'red'
    #     pyb.LED(2).on()

    # 处理STM32命令
    cmd_result = parse_command()

    if cmd_result['valid']:
        command_str = cmd_result['command']
        if command_str in OPERATION_MODE:
            current_operation = OPERATION_MODE[command_str]
            print(f"切换到{command_str}模式")

            # 同时处理需要颜色的模式
            if command_str in ['COLOR_DETECTION', 'COARSE_AREA_DETECTION']:
                target_color = cmd_result['color']
                pyb.LED(2).on()
            else:
                target_color = None
                pyb.LED(2).off()



    # 颜色检测模式且有指定目标颜色
    if current_operation == OPERATION_MODE['COLOR_DETECTION'] and target_color:
        # 执行特定颜色检测
        result = detect_specific_color(img, target_color)

        # 如果检测到颜色，发送结果
        if result:
            # 闪烁蓝色LED表示发送数据
            pyb.LED(3).toggle()

            # 发送颜色识别结果 (颜色结果代码, X坐标, Y坐标)
            send_data(result['code'], result['x'], result['y'])
            print(f"发送: 颜色={result['color']}, X={result['x']}, Y={result['y']}")



    # 修改后的COARSE_AREA_DETECTION处理
    elif current_operation == OPERATION_MODE['COARSE_AREA_DETECTION']:
        if target_color:
            threshold = COLOR_THRESHOLDS.get(target_color)
            if not threshold:
                continue

            # 颜色分割 + 形态学处理
            binary = img.binary(threshold).erode(1).dilate(1)

            # 查找最大色块
            blob = binary.find_blobs(area_threshold=5000)
            if blob:
                largest = max(blob, key=lambda b: b.area())
                cx, cy = largest.cx(), largest.cy()

                # 发送坐标
                norm_x = int(cx * 255 / img.width())
                norm_y = int(cy * 255 / img.height())
                send_data(norm_x, norm_y)
                print(norm_x, norm_y)


    # 显示帧率 (调试用)
    #print(f"FPS: {clock.fps():.1f}")
